import torch.nn as nn

#得到了输入之后直接根据cls进行分类
class BERT_SPC(nn.Module):
    def __init__(self, bert, opt):
        super(BERT_SPC, self).__init__()
        self.bert = bert
        self.dropout = nn.Dropout(opt.dropout)
        self.dense = nn.Linear(opt.bert_dim, opt.polarities_dim)  #只在bert基础上加了线性层？

    def forward(self, inputs):
        text_bert_indices, bert_segments_ids = inputs[0], inputs[1]

        #token_type_ids就是 token 对应的句子id，值为0或1（0表示对应的token属于第一句，1表示属于第二句）。
        '''
            _是batch*len*embedding size
            pooled_output:用于cls
        '''
        _, pooled_output = self.bert(text_bert_indices, token_type_ids=bert_segments_ids)  #目的是接到第二个输出
        pooled_output = self.dropout(pooled_output)
        logits = self.dense(pooled_output)
        return logits
